# -*- coding: utf-8 -*-
"""
@project: CloudDetection
@author: Wu Yue
@file: unet
@ide: PyCharm
@creat time: 2020-09-30 11:06
@change time:
@function: unet 网络
"""

import torch
# import torch.nn as nn
import torch.nn.modules as nn
import torch.nn.functional as F


"""
左半部分下采样路径组件
"""
class doubleConv(nn.Module):
    """
    横向两次卷积过程
    """
    def __init__(self, inCh: int, outCh: int):
        """

        :param inCh: 输入通道数
        :param outCh: 输出通道数
        """
        super(doubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=inCh, out_channels=outCh, kernel_size=3, padding=1),
            nn.BatchNorm2d(outCh),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=outCh, out_channels=outCh, kernel_size=3, padding=1),
            nn.BatchNorm2d(outCh),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class inConv(nn.Module):
    """
    下采样部分第一层卷积过程
    """
    def __init__(self, inCh: int, outCh: int):
        """

        :param inCh: 输入通道数
        :param outCh: 输出通道数
        """
        super(inConv, self).__init__()
        self.conv = doubleConv(inCh, outCh)

    def forward(self, x):
        x = self.conv(x)
        return x


class down(nn.Module):
    """
    一次下采样并完成横向卷积的过程
    """
    def __init__(self, inCh: int, outCh: int):
        """

        :param inCh: 输入通道数
        :param outCh: 输出通道数
        """
        super(down, self).__init__()
        self.mpConv = nn.Sequential(
            nn.MaxPool2d(kernel_size=2),
            doubleConv(inCh, outCh)
        )

    def forward(self, x):
        x = self.mpConv(x)
        return x


"""
右半部分上采样路径组件
"""
class up(nn.Module):
    """
    一次上采样并完成横向卷积的过程
    """
    def __init__(self, inCh: int, outCh: int, bilinear=True):
        """

        :param inCh: 输入通道数
        :param outCh: 输出通道数
        :param bilinear: 是否使用双线性插值
        """
        super(up, self).__init__()
        if bilinear:  # 声明使用的上采样方法为bilinear——双线性插值，默认使用这个值，计算方法为 floor(H*scale_factor)，所以由4*4变为8*8
            self.up = nn.Sequential(
                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
                nn.Conv2d(in_channels=inCh, out_channels=outCh, kernel_size=2, padding=0)
            )
        else:  # 否则就使用转置卷积来实现上采样，计算式子为 （Height-1）*stride - 2*padding -kernal_size +output_padding
            self.up = nn.Sequential(
                nn.ConvTranspose2d(inCh // 2, inCh // 2, 2, stride=2),
                nn.Conv2d(in_channels=inCh, out_channels=outCh, kernel_size=2, padding=0)
            )

        self.conv = doubleConv(inCh, outCh)

    def forward(self, x1, x2):
        """

        :param x1: 上采样传来的值
        :param x2: 下采样对应的值
        :return:
        """
        x1 = self.up(x1)
        x1 = F.pad(x1, [1, 0, 1, 0])
        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x


"""
Unet 网络整合
"""
class Unet(nn.Module):
    """
    Unet 模型
    """
    def __init__(self, nChannels, nClasses=1): #图片的通道数，1为灰度图像，3为彩色图像
        super(Unet, self).__init__()
        self.inputLayer = inConv(nChannels, 64)
        self.down1 = down(64, 128)
        self.down2 = down(128, 256)
        self.down3 = down(256, 512)
        self.down4 = down(512, 1024)
        self.up1 = up(1024, 512)
        self.up2 = up(512, 256)
        self.up3 = up(256, 128)
        self.up4 = up(128, 64)
        self.outputConv = nn.Conv2d(in_channels=64, out_channels=nClasses, kernel_size=1, padding=0)

    def forward(self, x):
        x1 = self.inputLayer(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)

        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        x = self.outputConv(x)

        return torch.sigmoid(x)  # 对预测结果进行二分类